﻿using System;
using System.IO;
using System.Net;
using System.Net.Sockets;
using System.Threading;

namespace woaC2.SocketLib
{
    /// <summary>
    /// Listen for policy file requests.
    /// </summary>
    public class PolicyServer
    {
        private const string PolicyRequestString = "<policy-file-request/>";

        private readonly byte[] _policy;

        private bool _isStopped;

        private TcpListener _listener;

        public PolicyServer(string policyRelativeFilePath)
        {
            var policyStream =
                new FileStream(
                    Path.Combine(AppDomain.CurrentDomain.BaseDirectory, policyRelativeFilePath), FileMode.Open);
            this._policy = new byte[policyStream.Length];
            policyStream.Read(this._policy, 0, this._policy.Length);
            policyStream.Close();
        }

        public void Start()
        {
            try
            {
                if (this._listener == null)
                {
                    this._listener = new TcpListener(IPAddress.Any, 943);
                }
                this._listener.Start();

                this._listener.BeginAcceptTcpClient(this.OnAcceptTcpClient, null);
                this._isStopped = false;
            }
            catch(Exception ex)
            {
                SocketLog.Write("[PolicyServer.Start]--[ERROR]--" + ex);
            }
        }

        public void Stop()
        {
            this._isStopped = true;

            try
            {
                if (_listener != null)
                    this._listener.Stop();
            }
            catch (Exception ex)
            {
                SocketLog.Write("[PolicyServer.Stop]--[ERROR]--" + ex);
            }
        }

        /// <summary>
        /// </summary>
        /// <param name="ar">
        /// </param>
        private void OnAcceptTcpClient(IAsyncResult ar)
        {
            if (this._isStopped)
                return;
            this._listener.BeginAcceptTcpClient(this.OnAcceptTcpClient, null);
            try
            {
                TcpClient client = this._listener.EndAcceptTcpClient(ar);
                ThreadPool.QueueUserWorkItem(o => this.SendPolicyFile(client));
            }
            catch (Exception ex)
            {
                SocketLog.Write("[PolicyServer.OnAcceptTcpClient]--[ERROR]--" + ex);
            }
        }

        private void SendPolicyFile(TcpClient client)
        {
            try
            {
                Stream clientStream = client.GetStream();
                var buffer = new byte[PolicyRequestString.Length];
                client.ReceiveTimeout = 5000;
                clientStream.Read(buffer, 0, buffer.Length);
                clientStream.Write(this._policy, 0, this._policy.Length);
                clientStream.Close();
                client.Close();
            }
            catch(Exception ex)
            {
                SocketLog.Write("[PolicyServer.SendPolicyFile]--[ERROR]--" + ex);
            }
        }
    }
}